from common.flask_package import *
import cv2
import numpy as np
from keras.models import load_model
from concurrent.futures import ThreadPoolExecutor
from ai.plate import place_str_relust
from common.my_fun import *
import requests
from common.my_msg import *


bp_video = Blueprint('video', __name__, url_prefix='/video')

# 加载训练好的模型
def load():
    global model
    model = load_model(r'ai/models/unet.h5')

# 打开视频文件
video = cv2.VideoCapture(0)
load()

# 调整图片大小
VIDEO_WIDTH = 640
VIDEO_HEIGHT = 480



# 视频流
def generate():
    with ThreadPoolExecutor(max_workers=2) as executor:
        while True:
            # 读取视频帧
            _, frame = video.read()
            if frame is None:
                break

            frame = cv2.resize(frame, (VIDEO_WIDTH, VIDEO_HEIGHT))
            # # 水平翻转图像并调整大小以适应html页面
            # frame = cv2.resize(cv2.flip(frame, 1), (960, 540))
            f = executor.submit(infer, frame)

            if f.result() is not None:
                confidence, x, y, w, h = f.result()
                # 在原图上标注车牌位置和置信度
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                text = f"Confidence: {confidence:.2f}"
                cv2.putText(frame, text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                # 置信度大于指定值执行逻辑
                if confidence >= 60:
                    cv2.imwrite("static/plateimg/result.jpg", frame) #保存结果图片
                    # 模拟发起请求
                    url = 'http://localhost:5000/car_info/add'
                    response = requests.post(url)
                    break

            # 将图片转换为字节流
            ret, jpeg = cv2.imencode('.jpg', frame)
            frame_bytes = jpeg.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n\r\n')


def infer(frame):
    # 对当前帧进行预测
    frame_resized = cv2.resize(frame, (512, 512))
    output_mask = model.predict(np.array([frame_resized]))[0]

    # 根据预测结果，在原图上标注出车牌的位置
    output_mask = cv2.resize(output_mask, (frame.shape[1], frame.shape[0]))
    output_mask_thresh = cv2.threshold(output_mask, 0.5, 255, cv2.THRESH_BINARY)[1].astype('uint8')

    # 将output_mask_thresh转换为CV_8UC1类型
    output_mask_thresh = cv2.cvtColor(output_mask_thresh, cv2.COLOR_BGR2GRAY)
    contours, hierarchy = cv2.findContours(output_mask_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) > 0:
        x, y, w, h = cv2.boundingRect(contours[0])
        # 在边界框上添加置信度文本
        confidence = model.predict(np.array([frame_resized]))[0][y:y + h, x:x + w].mean()
        return confidence, x, y, w, h
    return None

# 响应字节流
@bp_video.route('/video_feed')
def video_feed():
    return Response(generate(),mimetype='multipart/x-mixed-replace; boundary=frame')


